package net.maxupeng.asm74.segment;

import net.maxupeng.asm74.antlr.ASM74Parser;
import net.maxupeng.asm74.util.AntlrUtil;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static net.maxupeng.asm74.util.NumberUtil.str2int;


public class DataSegment extends Segment {
    public final Map<Integer, String> romLabelRefMap = new HashMap<>();
    public final Map<String, Integer> romLabelSizeMap = new HashMap<>();
    public ProcSegment initProc = null;


    public DataSegment(ASM74Parser.ProgContext progCtx) {
        super(progCtx);
    }

    @Override
    public void output_rom(int[] romArray) {
        super.output_rom(romArray);
        if (initProc != null) {
            initProc.output_rom(romArray);
        }
    }

    public void genInitProc() {
        int counter = 0;
        String code = ".segment \"CODE\"\n";
        code += ".proc _init_data" + getId() + "\n";

        for (Map.Entry<String, Integer> entry : romLabelSizeMap.entrySet()) {
            String label = entry.getKey();
            int size = entry.getValue();

            if (size <= 16) {
                for (int i = 0; i < size; i++) {
                    code += "                        lda _init_" + label + ",#" + i + "\n";
                    code += "                        sta " + label + ",#" + i + "\n\n";
                }
            } else {
                String label0 = "_init_data" + getId() + "_" + counter++;
                String label1 = "_init_data" + getId() + "_" + counter++;
                code += "                    lda #" + (size & 0xff) + "\n";
                code += "                    sta TMP0\n";
                code += "                    lda #" + ((size >> 8) & 0xff) + "\n";
                code += "                    sta TMP0,#1\n";
                code += "                    lda #L(_init_" + label + ")\n";
                code += "                    sta TMP1\n";
                code += "                    lda #H(_init_" + label + ")\n";
                code += "                    sta TMP1,#1\n";
                code += "                    lda #L(" + label + ")\n";
                code += "                    sta TMP2\n";
                code += "                    lda #H(" + label + ")\n";
                code += "                    sta TMP2,#1\n";
                code += "           " + label0 + ": lda TMP0\n";
                code += "                    or TMP0,#1\n";
                code += "                    jz " + label1 + "\n";
                code += "                    lda (TMP1)\n";
                code += "                    sta (TMP2)\n";
                code += "                    sec\n";
                code += "                    lda TMP0\n";
                code += "                    sub #1\n";
                code += "                    sta TMP0\n";
                code += "                    lda TMP0,#1\n";
                code += "                    sub #0\n";
                code += "                    sta TMP0,#1\n";
                code += "                    clc\n";
                code += "                    lda TMP1\n";
                code += "                    add #1\n";
                code += "                    sta TMP1\n";
                code += "                    lda TMP1,#1\n";
                code += "                    add #0\n";
                code += "                    sta TMP1,#1\n";
                code += "                    clc\n";
                code += "                    lda TMP2\n";
                code += "                    add #1\n";
                code += "                    sta TMP2\n";
                code += "                    lda TMP2,#1\n";
                code += "                    add #0\n";
                code += "                    sta TMP2,#1\n";
                code += "                    jmp " + label0 + "\n";
                code += "           " + label1 + ":\n\n";
            }
        }

        this.initProc = ProcSegment.parse(List.of(AntlrUtil.parse(code))).get(0);
        this.initProc.progCtx = this.progCtx;
    }

    public ProcSegment getInitProc() {
        return initProc;
    }

    public static List<DataSegment> parse(List<ASM74Parser.ProgContext> progCtxs) {
        List<DataSegment> datas = new ArrayList<>();

        DataSegment data = new DataSegment(null);
        String label = "";
        for (ASM74Parser.ProgContext progCtx : progCtxs) {
            boolean isDataSegment = false;
            for (ASM74Parser.LineContext lineCtx : progCtx.line()) {
                if (isDataSegment && lineCtx.label() != null) {
                    label = lineCtx.label().ID().getText();
                    data.romLabelOffsetMap.put("_init_" + label, data.romSize);
                    data.ramLabelOffsetMap.put(label, data.ramSize);
                    data.romLabelSizeMap.put(label, 0);
                }

                if (lineCtx.cmd() != null) {
                    ASM74Parser.CmdContext cmdCtx = lineCtx.cmd();
                    String cmd = cmdCtx.getStart().getText();
                    if (cmd.equals(".segment")) {
                        isDataSegment = cmdCtx.STRING().getText().trim().replaceAll("^\"|\"$", "").toUpperCase().equals("DATA");
                        if (isDataSegment) {
                            data = new DataSegment(progCtx);
                            datas.add(data);
                        }
                    }

                    if (isDataSegment) {
                        switch (cmd) {
                            case ".byte": {
                                List<Integer> bytes = new ArrayList<>();
                                for (int i = 0; i < cmdCtx.NUMBER().size(); i++) {
                                    bytes.add(str2int(cmdCtx.NUMBER(i).getText()));
                                }
                                for (Integer b : bytes) {
                                    data.rom.add(b & 0xff);
                                }
                                data.romSize += bytes.size();
                                data.ramSize += bytes.size();
                                data.romLabelSizeMap.put(label, data.romLabelSizeMap.get(label) + bytes.size());
                                break;
                            }
                            case ".word": {
                                List<Integer> words = new ArrayList<>();
                                for (int i = 0; i < cmdCtx.NUMBER().size(); i++) {
                                    words.add(str2int(cmdCtx.NUMBER(i).getText()));
                                }
                                for (Integer w : words) {
                                    data.rom.add(w & 0xff);
                                    data.rom.add((w >> 8) & 0xff);
                                }
                                data.romSize += words.size() * 2;
                                data.ramSize += words.size() * 2;
                                data.romLabelSizeMap.put(label, data.romLabelSizeMap.get(label) + words.size() * 2);
                                break;
                            }
                            case ".dword": {
                                List<Integer> dwords = new ArrayList<>();
                                for (int i = 0; i < cmdCtx.NUMBER().size(); i++) {
                                    dwords.add(str2int(cmdCtx.NUMBER(i).getText()));
                                }
                                for (Integer dw : dwords) {
                                    data.rom.add(dw & 0xff);
                                    data.rom.add((dw >> 8) & 0xff);
                                    data.rom.add((dw >> 16) & 0xff);
                                    data.rom.add((dw >> 24) & 0xff);
                                }
                                data.romSize += dwords.size() * 4;
                                data.ramSize += dwords.size() * 4;
                                data.romLabelSizeMap.put(label, data.romLabelSizeMap.get(label) + dwords.size() * 4);
                                break;
                            }
                            case ".addr": {
                                // TODO 这里没处理函数指针的情况
                                data.romLabelRefMap.put(data.romSize, cmdCtx.ID().getText());
                                data.ramSize += 2;
                                data.romSize += 2;
                                data.rom.add(0);
                                data.rom.add(0);
                                data.romLabelSizeMap.put(label, data.romLabelSizeMap.get(label) + 2);
                                break;
                            }
                            case ".res": {
                                int size = str2int(lineCtx.cmd().NUMBER(0).getText());
                                data.ramSize += size;
                                break;
                            }
                        }
                    }
                }
            }
        }

        return datas;
    }
}
